#include "opencv2\opencv.hpp"
#include "cmd_utility.h"
#include "facade_grammar.h"
#include "facade_model.h"
#include "svm.h"

#include <iostream>
#include <vector>


CvRect box;
bool drawing_box = false;
int n_terminals = 0;
int current_terminal = 0;
CvScalar colors[] = {CV_RGB(255, 0, 0), CV_RGB(0, 255, 0), CV_RGB(0, 0, 255), CV_RGB(255, 255, 0),
	CV_RGB(255, 0, 255), CV_RGB(0, 255, 255)};

CommandLineArgument cmdline_arg;
std::shared_ptr<FacadeGrammar> facade_grammar;
std::shared_ptr<FacadeModel> facade_model;


svm_problem problem;
svm_parameter parameter;
svm_model* model;

std::vector<std::vector<CvPoint> > training_samples;

void train();
void predict();
double grid_search(svm_problem* prob, svm_parameter* param, int nr_fold);

// A little subroutine to draw a box onto an image
void draw_box(IplImage* img, CvRect rect);

// Define our callback which we will install for mouse events
void my_mouse_callback(int event, int x, int y, int flags, void* param);

int main(int argc, char** argv) {
	box = cvRect(-1, -1, 0, 0);

	// Parsing the command line
	if (!cmdline_arg.ParseCommandLineArgument(argc, argv)) {
		std::cout << "parsing command line fail!\n";
		exit(-1);
	}

	// Load the grammar
	facade_grammar = std::shared_ptr<FacadeGrammar>(new FacadeGrammar);
	assert(facade_grammar->ReadGrammar(cmdline_arg.grammar_));
	n_terminals = facade_grammar->get_terminal_symbol_size();
	training_samples.resize(n_terminals);

	// Load the image
	std::ostringstream filename;
	filename << cmdline_arg.grid_ << "depth" << ".png";
	IplImage* img = cvLoadImage(filename.str().c_str());
	if (img == NULL) {
		std::cout << "Load image fail!\n";
		exit(-1);
	}
	// backup  the image
	IplImage* tmp = cvCloneImage(img);

	// Create a named window
	cvNamedWindow("SVM-Classifier");

	// Install the call back
	cvSetMouseCallback("SVM-Classifier", my_mouse_callback, (void*)img);

	// The main program loop.
	// Here we copy the working image to the 'temp' image,
	// and if the user is drawing, then put the currently
	// contemplated box onto that temp image. display the temp
	// image, and wait 15ms for a key stroke, then repeat ...
	while(1) {
		if (tmp != NULL)
			cvReleaseImage(&tmp);
		tmp = cvCloneImage(img);
		if (drawing_box) 
			draw_box(tmp, box);
		cvShowImage("SVM-Classifier", tmp);

		int c = cvWaitKey(15);

		if ((c & 255)== 27) {
			std::cout << "Exiting ...\n";
			break;
		}
		switch ((char)c)
		{
		case 'c':
			++current_terminal;
			if (current_terminal >= n_terminals) {
				current_terminal = 0;
			}
			std::cout << "current terminal index: " << current_terminal << std::endl;
			break;
		case 't':
			// load the attributes
			facade_model = std::shared_ptr<FacadeModel>(new FacadeModel(facade_grammar));
			facade_model->LoadGrids(cmdline_arg.grid_);

			if (!cmdline_arg.load_svm_model_) {
				train();
			} else {
				model = svm_load_model(cmdline_arg.svm_model_file_.c_str());
				std::cout << "SVM model loaded.\n";
			}
			break;
		case 'p':
			predict();
			facade_model->RenderGrd(cmdline_arg.grid_);
			facade_model->SaveGrids(cmdline_arg.grid_);
			std::cout << "Classify results saved.\n";
			break;
		default:
			break;
		}
	}

	// Be tidy
	cvReleaseImage(&img);
	cvReleaseImage(&tmp);
	cvDestroyWindow("SVM-Classifier");

	return 0;
}

// This is our mouse callback. If the user
// presses the left button, we start a box.
// when the under releases that button, then we
// add the box to he current image. When the mouse is
// dragged(width the button down) we resize the box
void my_mouse_callback(int event, int x, int y, int flags, void* param) {
	IplImage* image = (IplImage*)param;

	switch (event)
	{
	case CV_EVENT_MOUSEMOVE:
		if (drawing_box) {
			box.width = x - box.x;
			box.height = y - box.y;
		}
		break;
	case CV_EVENT_LBUTTONDOWN:
		drawing_box = true;
		box = cvRect(x, y, 0, 0);
		break;
	case CV_EVENT_LBUTTONUP:
		drawing_box = false;
		if (box.width < 0) {
			box.x += box.width;
			box.width *= -1;
		}
		if (box.height < 0) {
			box.y += box.height;
			box.height *= -1;
		}
		draw_box(image, box);

		for (int i = box.x; i < box.x + box.width; ++i) {
			for (int j = box.y; j < box.y + box.height; ++j) {			
				training_samples[current_terminal].push_back(cvPoint(i, j));
			}
		}

		std::cout << "terminal " << current_terminal << " has " << training_samples[current_terminal].size() << " pixels.\n";
		break;
	default:
		break;
	}
}

void draw_box(IplImage* img, CvRect rect) {
	cvRectangle(img, cvPoint(box.x, box.y),
		cvPoint(box.x + box.width, box.y + box.height),
		colors[current_terminal],
		CV_FILLED
		);
}

void train() {
	/* Construct SVM problem */
	std::cout << "Construct SVM problem ...\n";
	problem.l = 0;
	for (int i = 0; i < n_terminals; ++i) {
		problem.l += training_samples[i].size();
	}
	problem.y = new double[problem.l];
	problem.x = new svm_node*[problem.l];

	int label_count = 0; 
	int sample_count = 0;
	for (int i = 0; i < n_terminals; ++i) {
		std::vector<CvPoint>& samples = training_samples[i];
		for (int j = 0; j < samples.size(); ++j) {
			// label
			problem.y[label_count++] = i;
			// features
			problem.x[sample_count] = new svm_node[3];

			problem.x[sample_count][0].index = 1;
			problem.x[sample_count][0].value = facade_model->GetProperty(samples[j].x, samples[j].y, 0);
			problem.x[sample_count][1].index = 2;
			problem.x[sample_count][1].value = facade_model->GetProperty(samples[j].x, samples[j].y, 1);
			problem.x[sample_count++][2].index = -1;

		}
	}
	assert(label_count == problem.l);
	assert(sample_count == problem.l);

	/* Specify SVM parameters */
	std::cout << "Specify SVM parameters ...\n";
	// default values
	svm_parameter parameter;
	parameter.svm_type = C_SVC;
	parameter.kernel_type = RBF;
	parameter.degree = 3;
	parameter.gamma = 1;	// 1/num_features
	parameter.coef0 = 0;
	parameter.nu = 0.5;
	parameter.cache_size = 100;
	parameter.C = 1;
	parameter.eps = 1e-3;
	parameter.p = 0.1;
	parameter.shrinking = 0;
	parameter.probability = 1;
	parameter.nr_weight = 0;
	parameter.weight_label = NULL;
	parameter.weight = NULL;

	// check the parameters
	const char* error_infor = svm_check_parameter(&problem, &parameter);
	if (error_infor != NULL){
		std::cout << error_infor << std::endl;
		exit(-1);
	}

	/* Cross Validate */
	std::cout << "Cross Validate ...\n";
	double accuracy = grid_search(&problem, &parameter, n_terminals);
	std::cout << "best c: " << parameter.C << " best g: " << parameter.gamma << " Accuracy: " << accuracy << std::endl;

	/* Training */
	std::cout << "Training\n";
	model = svm_train(&problem, &parameter);

	svm_save_model(cmdline_arg.svm_model_file_.c_str(), model);

	std::cout << "SVM model saved.\n";
}

void predict() {
	std::cout << "Predict ...\n";
	for ( int i = 0; i < facade_model->get_width(); ++i) {
		for (int j = 0; j < facade_model->get_height(); ++j) {
			svm_node* node;
			node = new svm_node[3];

			node[0].index = 1;
			node[0].value = facade_model->GetProperty(i, j, 0);
			node[1].index = 2;
			node[1].value = facade_model->GetProperty(i, j, 1);
			node[2].index = -1;

			// double label = svm_predict(model, node);
			const int nr_class = svm_get_nr_class(model);
			assert(nr_class == n_terminals);

			double* probility_estimates = new double[nr_class];

			double label = svm_predict_probability(model, node, probility_estimates);
		
			facade_model->SetProperty(i, j, 2, label);

			for (int k = 0; k < nr_class; ++k) {
				facade_model->SetProperty(i, j, k + 3, probility_estimates[k]);
			}

			delete[] probility_estimates;
		}
	}
	std::cout << "Done.\n";
}

double grid_search(svm_problem* prob, svm_parameter* param, int nr_fold) {
	int log2c_begin = -5;
	int log2c_end = 15;
	int log2c_step = 2;
	int log2g_begin = 3;
	int log2g_end = -15;
	int log2g_step = -2;

	double* target = new double[prob->l];

	int best_correct = 0;
	double best_c = 1;
	double best_g = 1;
	double best_accuracy;

	for (int log2c = log2c_begin; log2c <= log2c_end; log2c += log2c_step) {
		for (int log2g = log2g_begin; log2g >= log2g_end; log2g += log2g_step) {
			param->C = pow(2.0, log2c);
			param->gamma = pow(2.0, log2g);

			svm_cross_validation(prob, param, nr_fold, target);
			int total_correct = 0;
			for(int i = 0; i  < prob->l; ++i)
				if(target[i] == prob->y[i])
					++total_correct;
			if (total_correct >= best_correct) {
				best_correct = total_correct;
				best_c = param->C;
				best_g = param->gamma;
			}
		}
	}

	param->C = best_c;
	param->gamma = best_g;

	delete[] target;
	best_accuracy = 100.0 * best_correct / prob->l;
	return best_accuracy;
}